{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "YYZqhjGhK8nJ"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "V1-OvloqK4CX"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VKRNwMTxmZXx"
      },
      "source": [
        "# Concrete functions\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/guide/concrete_function\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003e\n",
        "    View on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/concrete_function.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003e\n",
        "    Run in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/concrete_function.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003e\n",
        "    View source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/concrete_function.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VgQvKIvXockG"
      },
      "source": [
        "In the guide to [AutoGraph and `tf.functions`](function.ipynb) you saw how to use `tf.function`. This guide dives into the details of: \n",
        "\n",
        "* `tf.function` Tracing\n",
        "* `tf.function` Signatures\n",
        "* The Concrete functions generated by tracing:\n",
        "  * How to access them\n",
        "  * How to use them\n",
        "\n",
        "These details only become important:\n",
        "\n",
        "* If you're experiencing performance issues due to undesired tracing\n",
        "  of a `tf.funcion`.\n",
        "* When you need precise control over the TensorFlow Graphs generated by\n",
        "  `tf.function`. For example for exporting the model to\n",
        "  [TFLite](https://tensorflow.org/lite/converter) using\n",
        "  `tf.lite.Converter.from_concrete_functions`.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "LwLkhiXIrT5m"
      },
      "source": [
        "## Background\n",
        "\n",
        "In TensorFlow 2, eager execution is on by default. TensorFlow's eager\n",
        "execution is an imperative programming environment that evaluates operations\n",
        "immediately, without building graphs. Operations return values instead\n",
        "of constructing a computational graph to run later. Here is a [detailed guide on eager\n",
        "execution](eager.ipynb).\n",
        "\n",
        "Running imperatively makes development and debugging\n",
        "more interactive, but doesn't allow for easy exporting.\n",
        "\n",
        "The `tf.function` API makes it possible to save models as graphs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "GgJHWoOPrWtK"
      },
      "source": [
        "## Terminology\n",
        "\n",
        "The following terminology is used in this document:\n",
        "\n",
        "*   **Signature** - A description of the inputs and outputs for a set of operations.\n",
        "* **Polymorphic function** - Python callable that encapsulates several\n",
        "    concrete function graphs behind one API.  \n",
        "* **Concrete function** - Graph with a single signature.\n",
        "  \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "vq126xVZr2CL"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "jtRSqqP6o_pQ"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import, division, print_function, unicode_literals"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3_jlx1_ir1hb"
      },
      "outputs": [],
      "source": [
        "import traceback\n",
        "import textwrap\n",
        "\n",
        "try:\n",
        "  !pip install tf-nightly\n",
        "except Exception:\n",
        "  pass"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "AB6fMBGSofmm"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TpkKbIBerbWi"
      },
      "source": [
        "## Create a `tf.function`\n",
        "\n",
        "Annotating a function with `tf.function` generates a *polymorphic function*\n",
        "containing those operations. All operations that are not annotated with\n",
        "`tf.function` will be evaluated with eager execution. The examples below show a quick example of `tf.function` usage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "fPxeVDQ2mZXy"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def square(x):\n",
        "  return x*x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "yDO0s_LWBc9S"
      },
      "outputs": [],
      "source": [
        "square(2).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8AhtB2hAsByK"
      },
      "source": [
        "Remember that the python decorator syntax just calls the decorator with the decorated object as input:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "p-KckIFcmZX2"
      },
      "outputs": [],
      "source": [
        "def pow(x,y):\n",
        "  return x ** y\n",
        "\n",
        "pow = tf.function(pow)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1moLFTi2sToR"
      },
      "outputs": [],
      "source": [
        "pow(3,4).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hMhMhl3lmZX6"
      },
      "source": [
        "### Attach a `tf.function` method to a `tf.Module`\n",
        "\n",
        "The `tf.function` can be optionally stored as part of a `tf.Module` object. The `tf.Module` class provides features for tracking variables and saving [checkpoints](checkpoints.ipynb) and [models](saved_models).\n",
        "\n",
        "Classes like `keras.layers.Layer` and `keras.Model` are subclasses of Module."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ht2K2vpDmZX7"
      },
      "outputs": [],
      "source": [
        "class Pow(tf.Module):\n",
        "  def __init__(self, exponent):\n",
        "    self.exponent = tf.Variable(exponent, dtype = tf.float32, name='Pow/exponent')\n",
        "\n",
        "  @tf.function\n",
        "  def __call__(self, x):\n",
        "    return x ** self.exponent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1GVBiQfx8klb"
      },
      "outputs": [],
      "source": [
        "pow = Pow(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "KqnKlZe58nPe"
      },
      "outputs": [],
      "source": [
        "pow.variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RQBVvuxLtGz2"
      },
      "outputs": [],
      "source": [
        "pow(tf.constant(2.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "kKiJrZlS5eeT"
      },
      "outputs": [],
      "source": [
        "pow.exponent.assign(4)\n",
        "pow(tf.constant(2.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "t_HlUmTe5reS"
      },
      "outputs": [],
      "source": [
        "tf.saved_model.save(pow, 'pow')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "abpBoMRR5ypu"
      },
      "outputs": [],
      "source": [
        "reloaded_pow = tf.saved_model.load('pow')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "GA_a3CH956MR"
      },
      "outputs": [],
      "source": [
        "reloaded_pow(tf.constant(3.0)).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "rggP9vCK60R4"
      },
      "source": [
        "### Assign a `tf.function` as an attribute\n",
        "If you assign a `tf.Module` or a `tf.function` as an attribute of a module it will be serialized as well:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "kR2-Mwyg7GyO"
      },
      "outputs": [],
      "source": [
        "mod = tf.Module()\n",
        "mod.increment_by = tf.Variable(2.0)\n",
        "\n",
        "@tf.function\n",
        "def increment(x):\n",
        "  return x+mod.increment_by\n",
        "\n",
        "mod.inc = increment\n",
        "mod.inc(tf.constant(1.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "rR5D4q7Z76Uo"
      },
      "outputs": [],
      "source": [
        "mod.cube = Pow(3)\n",
        "mod.cube(tf.constant(2.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "vNIzKT8I8XzC"
      },
      "outputs": [],
      "source": [
        "mod.variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "OQqnUgjz7QOW"
      },
      "outputs": [],
      "source": [
        "tf.saved_model.save(mod, 'mod')\n",
        "reloaded_mod = tf.saved_model.load('mod')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Uvikzz7x7jB1"
      },
      "outputs": [],
      "source": [
        "reloaded_mod.inc(4.0).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "n7eVyFhT8Bq6"
      },
      "outputs": [],
      "source": [
        "reloaded_mod.cube(4.0).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Sj7P9wBFmZYW"
      },
      "source": [
        "### Interoperability with `tf.keras`\n",
        "\n",
        "Keras classes like `keras.Model` and `keras.layers.Layer` are fully compatible with `tf.function` and `tf.Module`.\n",
        "\n",
        "For example, build a simple model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "WYHD6NMWmZYX"
      },
      "outputs": [],
      "source": [
        "linear = tf.keras.Sequential([tf.keras.layers.Dense(units=1, input_shape=[1])])\n",
        "linear.compile(optimizer='adam', loss='mean_squared_error')\n",
        "linear.fit(x=[-1, 0, 1, 2, 3, 4], y=[-3, -1, 1, 3, 5, 7], epochs=50, verbose=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "5rumHqek5oMN"
      },
      "outputs": [],
      "source": [
        "linear(tf.constant([[1],[2]]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "dx2HIyx_UI9N"
      },
      "source": [
        "Inspect it's variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "vLpavpL-86kn"
      },
      "outputs": [],
      "source": [
        "linear.variables"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JEGvkaGH5hIZ"
      },
      "source": [
        "Now attach it to a `tf.Module`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "muWIMr6-5B8A"
      },
      "outputs": [],
      "source": [
        "module = tf.Module()\n",
        "module.linear = linear"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SF24uHF59kH4"
      },
      "source": [
        "The `tf.Module` also tracks the `tf.Variable`s:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Z6ovPDgT5RYz"
      },
      "outputs": [],
      "source": [
        "module.variables"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-ZyWZdJVtDp1"
      },
      "source": [
        "The `tf.Module` will export the contents of the `keras.Model` as well:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ZgNSWOoi6ioK"
      },
      "outputs": [],
      "source": [
        "tf.saved_model.save(module,'module')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RS3j67MK9zVR"
      },
      "outputs": [],
      "source": [
        "reloaded = tf.saved_model.load('module')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "pxDCTrZd965w"
      },
      "outputs": [],
      "source": [
        "reloaded.linear([[1.0]])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ROhab4u7idvG"
      },
      "source": [
        "## Tracing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zQaBkCHB-W0Q"
      },
      "source": [
        "The objects returned from `tf.function` are polymorphic functions. They will accept python objects, or `tf.Tensors` with any shape or `tf.dtype` as input.\n",
        "\n",
        "In the background TensorFlow builds `tf.Graph`s representing the calculation. \n",
        "This graph is wrapped in a python callable: a concrete function. Each concrete function can only handle a single input signature.\n",
        "\n",
        "`tf.function` traces the python function each time in needs to create a concrete function. The easiest way to see when a function is traced is to add a call to print: \n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "FgvSApDrmJGq"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def mul(a, b):\n",
        "  print('Tracing:\\n    {a}\\n    {b}\\n'.format(a=a, b=b))\n",
        "  return a*b"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "jxzhLeKV_chb"
      },
      "source": [
        "### Dtypes and shapes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "i_OnPSdOCI7J"
      },
      "source": [
        "If you call the polymorphic function with two different types of input, it will trace once for each:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "A2uKXbRMmTGj"
      },
      "outputs": [],
      "source": [
        "# Trace with ints\n",
        "mul(tf.constant(2), tf.constant(3)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "a_rcmkWxCVwE"
      },
      "outputs": [],
      "source": [
        "# Trace with floats\n",
        "mul(tf.constant(2.0), tf.constant(3.0)).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Jmqq9rrDCdI4"
      },
      "source": [
        "When you call it again with the same input types, it dispatches to an existing function instead of tracing:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "KFtAlPZPqM7A"
      },
      "outputs": [],
      "source": [
        "# Call with ints again =\u003e no trace\n",
        "mul(tf.constant(10), tf.constant(10))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "mH920uLTCwK1"
      },
      "source": [
        "Changing the sizes of the inputs also triggers a trace (setting `tf.function(experimental_relax_shapes=True)` may reduce this): "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "eb_UByibqGny"
      },
      "outputs": [],
      "source": [
        "# Trace with vectors\n",
        "mul(tf.constant([1.0,3.0]), tf.constant(3.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "eXm0fgVzqVxT"
      },
      "outputs": [],
      "source": [
        "# Trace with different-sized vectors\n",
        "mul(tf.constant([1.0,2.0,3.0, 4.0]), tf.constant(3.0))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "D-LoC287Ax5k"
      },
      "source": [
        "### Immutable python objects"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "RpkfdNeKFbiU"
      },
      "source": [
        "If you pass an immutable python object, like a `int`, `str`, or `tuple` to a `tf.function`, it executes a trace for each *value* of those python objects.\n",
        "\n",
        "This is useful to control what gets included in the `tf.Graph` (See: [The Autograph Guide](function.ipynb) for more details).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "-G1N7qh4BAhe"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def mul(a, b):\n",
        "  print('Tracing:\\n    {a}\\n    {b}\\n'.format(a=a, b=b))\n",
        "  return a*b"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ftCJz7HfDGof"
      },
      "outputs": [],
      "source": [
        "# Trace for a=3.0\n",
        "mul(3.0, tf.constant(3.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "C8hiJCP_Du0I"
      },
      "outputs": [],
      "source": [
        "# Don't trace for a=3.0 the second time:\n",
        "mul(3.0, tf.constant(3.0)).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Sm8Cq0uTt88j"
      },
      "source": [
        "Caution: It is easy to cause many traces by passing unique python values. This can be a significant performance problem. Often passing a `tf.Tensor` value is the solution.\n",
        "\n",
        "This loop traces the function for each unique int:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "4LTTGqQryjsj"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def power(a,b):\n",
        "  print('Tracing \"power\": a={}'.format(a))\n",
        "  return a**b"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "yyGenHr_t8LK"
      },
      "outputs": [],
      "source": [
        "p = tf.constant(2)\n",
        "for n in range(12):\n",
        "  power(n,p)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xYgJ9l7_yY20"
      },
      "source": [
        "On the second run each int has been traced, so there's no tracing to do:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "duPVVOg3yUCl"
      },
      "outputs": [],
      "source": [
        "p = tf.constant(2)\n",
        "for n in range(12):\n",
        "  power(n,p)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "6bX9b1t0uTQX"
      },
      "source": [
        "To avoid excess retracing be sure to pass a `tf.Tensor` instead of python numbers or strings:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "F78ik9psuPn7"
      },
      "outputs": [],
      "source": [
        "p = tf.constant(2)\n",
        "for n in tf.range(12):\n",
        "  power(n,p)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "N-wwjPtF1Q25"
      },
      "source": [
        "To shut off tracing altogether, pass a signature to the `tf.function` decorator:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "9EcEZfZY1Ylg"
      },
      "outputs": [],
      "source": [
        "@tf.function(input_signature=(\n",
        "    tf.TensorSpec(shape=[], dtype=tf.float32),\n",
        "    tf.TensorSpec(shape=[], dtype=tf.float32),)\n",
        ")\n",
        "def power_with_sig(a,b):\n",
        "  print('Tracing \"power_with_sig\"')\n",
        "  return a**b"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Pff46f4U1-6R"
      },
      "outputs": [],
      "source": [
        "power_with_sig(3.0, 3.0).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "cBT2IFiZ1wLS"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  power_with_sig(tf.constant([1.0,2.0,3.0]),tf.constant(3.0))\n",
        "  assert False\n",
        "except ValueError:\n",
        "  traceback.print_exc(limit=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kjeEQ4J4tt81"
      },
      "source": [
        "### Example: Dropout"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-QbejAfEDLNX"
      },
      "source": [
        "Retracing for specific values gives you control over what code gets generated by the `tf.function`.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "kIgq7vO1sXh7"
      },
      "outputs": [],
      "source": [
        "class Dropout(tf.Module):\n",
        "  def __init__(self, rate, name=None):\n",
        "    super(Dropout, self).__init__(name)\n",
        "    self.rate = tf.Variable(rate, dtype = tf.float32, trainable=False)\n",
        "\n",
        "  @tf.function\n",
        "  def __call__(self, x, training=True):\n",
        "    print(textwrap.dedent(\"\"\"\n",
        "                          Tracing \"Dropout\":\n",
        "                              training = {}\n",
        "                              x = {}\n",
        "                              name = {:s}\n",
        "                          \"\"\".format(training, x, self.name)))\n",
        "    if training:\n",
        "      print('    - Train branch\\n')\n",
        "      mask = tf.random.uniform(x.shape) \u003e self.rate\n",
        "      return x * tf.cast(mask, tf.float32)/self.rate\n",
        "    else:\n",
        "      print('    - Test branch\\n')\n",
        "      return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "o7TFSeINqhtc"
      },
      "source": [
        "Create an instance of this simple `Dropout` layer:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "5DHE1io1pN2t"
      },
      "outputs": [],
      "source": [
        "dropout = Dropout(0.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "CE4fuvUP0BWE"
      },
      "source": [
        "The first time you call it with a python `training=True` as input, it traces the `training` branch:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "O5xIiGATqNQY"
      },
      "outputs": [],
      "source": [
        "dropout(tf.range(10, dtype=tf.float32), training=True).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Cup3pXHI0WWH"
      },
      "source": [
        "The second time, it doesn't need to re-trace the branch:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VaMDgyM8zpge"
      },
      "outputs": [],
      "source": [
        "dropout(tf.range(10, dtype=tf.float32), training=True).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yhiyPC7s0dF9"
      },
      "source": [
        "Passing `training=False` triggers a trace on the first run since this is a different python value:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mgaR5PFTzk9A"
      },
      "outputs": [],
      "source": [
        "dropout(tf.range(10, dtype=tf.float32), training=False).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "DVeAfMkbzn-V"
      },
      "outputs": [],
      "source": [
        "dropout(tf.range(10, dtype=tf.float32), training=False).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "G6b57DUG32b7"
      },
      "source": [
        "If you pass a `bool` tensor, it uses TensorFlow autograph rewrite the `if` to a `tf.cond`m and traces both branches:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "s_FsCDEaqYqq"
      },
      "outputs": [],
      "source": [
        "dropout(tf.range(10, dtype=tf.float32), training=tf.constant(False)).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "UbmPUjhI8Rgf"
      },
      "source": [
        "This captures the control flow in a single concrete function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "BViHlbCN0wss"
      },
      "outputs": [],
      "source": [
        " dropout(tf.range(10, dtype=tf.float32), training=tf.constant(True)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "p8eT9zKXp8F5"
      },
      "outputs": [],
      "source": [
        "dropout(tf.range(10, dtype=tf.float32), training=tf.constant(False)).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "smyOmjl2GCdt"
      },
      "source": [
        "### Other python objects\n",
        "\n",
        "Since the generated `tf.Graphs` cannot contain complex python objects, these are included by tracing and variable capture. \n",
        "\n",
        "The `tf.function` runs a separate trace for each **instance**. So each trace includes its own variables, and can set its behavior based on the instance.\n",
        "\n",
        "The most common usage is on methods of Module, Layer or Module:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "C6SBvOEdO5Qw"
      },
      "outputs": [],
      "source": [
        "dropout_a = Dropout(0.5, name='dropout_a')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "uOC8COMxQY0J"
      },
      "outputs": [],
      "source": [
        "print(dropout_a(tf.range(10, dtype=tf.float32), True).numpy())\n",
        "print(dropout_a(tf.range(10, dtype=tf.float32), True).numpy())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mISom9-ZTf5g"
      },
      "outputs": [],
      "source": [
        "dropout_b = Dropout(0.5, name='dropout_b')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "95OE0xxeQ4t0"
      },
      "outputs": [],
      "source": [
        "print(dropout_b(tf.range(10, dtype=tf.float32), True).numpy())\n",
        "print(dropout_b(tf.range(10, dtype=tf.float32), True).numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Ts8nQVMcQfsK"
      },
      "source": [
        "But the behavior is the same on a stand-alone `tf.function`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "0xVE3zoLs61l"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def run(callable, x):\n",
        "  print('Tracing \"run\":\\n    callable = {}\\n    x = {}\\n'.format(callable, x))\n",
        "  return callable(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "WAs-suN2tjH0"
      },
      "outputs": [],
      "source": [
        "def plus_1(x):\n",
        "  return x+1\n",
        "\n",
        "print(run(plus_1, tf.constant(2.0)).numpy())\n",
        "print(run(plus_1, tf.constant(5.0)).numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "g_m9M64oXFeW"
      },
      "source": [
        "The tracing one `tf.function` can trigger tracing in another:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "lnib6kNTu4bF"
      },
      "outputs": [],
      "source": [
        "print(run(dropout, tf.range(10.0)).numpy())\n",
        "print(run(dropout, tf.range(10.0)).numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4DE76z8KiVTG"
      },
      "source": [
        "### Weak references"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "LKaSfW5likUT"
      },
      "source": [
        "Caution: Each trace only keeps a [weak-reference](https://docs.python.org/3/library/weakref.html) to any `tf.Variable`. If the variable is not kept alive by another reference, the trace may become unusable."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yGUY0lVD8nMT"
      },
      "source": [
        "For example here's a `tf.function` that refers to `var` from the enclosing scope:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "zwwnJwXR8PEb"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def plus_var(x):\n",
        "  print('Tracing \"plus_var\":\\n    x = {}\\n    var = {}\\n\\n'.format(x, var.name))\n",
        "  return x + var"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9qmxAmFl8v5p"
      },
      "source": [
        "Trace the function with one variable:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Nebd4e5y8UYa"
      },
      "outputs": [],
      "source": [
        "var = tf.Variable(1, name=\"IntVar\")\n",
        "plus_var(tf.constant([1,2])).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DaqblviJ8y6v"
      },
      "source": [
        "And with another variable:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "s3EwG-yt8Zgx"
      },
      "outputs": [],
      "source": [
        "var = tf.Variable(2.0, name=\"FloatVar\")\n",
        "plus_var(tf.constant([2.0, 10.0])).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "cJL43oZN86o4"
      },
      "source": [
        "That worked, but because you no longer have a reference to `\"IntVar\"`, that first trace is broken:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Yb4RkESA8gs2"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  plus_var(tf.constant([1,2])).numpy()\n",
        "  assert False\n",
        "except tf.errors.FailedPreconditionError:\n",
        "  traceback.print_exc(limit=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "235EbM2HmZYB"
      },
      "source": [
        "## Accessing concrete function\n",
        "\n",
        "In the previous section you saw the conditions for triggering a new trace of a polymorphic `tf.function`. Each trace generates a new concrete function.\n",
        "\n",
        "When you save `tf.Module` as a `tf.saved_model` It's those concrete functions that define the `tf.Graph`s that are exported. You don't save a `tf.function` you save the concrete functions that are created by tracing. \n",
        "\n",
        "To get a concrete function from the polymorphic `tf.function` you need to define the signature. Either:\n",
        "\n",
        "*   Pass an `input_signature` to `tf.function`, and call the \n",
        "    `get_concrete_function()` method.\n",
        "*   Pass a list of `tf.TensorSpec`s to `get_concrete_function`: `tf.TensorSpec(shape=[1], dtype=tf.float32)`.\n",
        "*   Pass an example tensor of the correct shape and type to\n",
        "    `get_concrete_function`: `tf.constant(1., shape=[1])`.\n",
        "\n",
        "The following example shows how to define the `input_signature` parameter for\n",
        "`tf.function`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "O-tE5Pn6h1DU"
      },
      "source": [
        "#### Using `input_signature`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hsMR5JBdhykq"
      },
      "source": [
        "Specify input tensors in the call to `tf.function` as shown below.\n",
        "This `tf.function`can only execute on tensors that match the specified signatutre.\n",
        "\n",
        "A `None` in the `shape` acts a wildcard. So this these `tf.TensroSpec` say \"A float32 vector of any length\".\n",
        "\n",
        "This pattern can be very important if your `tf.function` is expected to handle sequences of different length, or images of different sizes for each batch (See [Transformer](../tutorials/text/transformer.ipynb) and [Deep Dream](../tutorials/generative/deepdream.ipynb) tutrorials for example)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "INJhW9c5htAK"
      },
      "outputs": [],
      "source": [
        "@tf.function(input_signature=(\n",
        "    tf.TensorSpec(shape=[None], dtype=tf.float32),\n",
        "    tf.TensorSpec(shape=[None], dtype=tf.float32),)\n",
        ")\n",
        "def power_with_sig(a,b):\n",
        "  print('Tracing \"power_with_sig\"\\n')\n",
        "  return a**b"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Q0U5ytANi4V7"
      },
      "source": [
        "Calling `get_concrete_function` will execute the trace (if necessary), and return a concrete function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "yR3tqv1pijvh"
      },
      "outputs": [],
      "source": [
        "p = power_with_sig.get_concrete_function()\n",
        "type(p)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "XvGhi41EipeQ"
      },
      "outputs": [],
      "source": [
        "p(tf.constant([2.0,3.0,4.0]), tf.constant([5.0,4.0,3.0])).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "bBKfDEeMz43m"
      },
      "source": [
        "### Using `get_concrete_function`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "IGlmMUSK0c3H"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def power(a,b):\n",
        "  print('Tracing \"power\"\\n')\n",
        "  return a**b"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "NYbxGotFqoDB"
      },
      "outputs": [],
      "source": [
        "float_power = power.get_concrete_function(\n",
        "  a = tf.TensorSpec(shape=[], dtype=tf.float32),\n",
        "  b = tf.TensorSpec(shape=[], dtype=tf.float32))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1cXOt1g0qpCL"
      },
      "outputs": [],
      "source": [
        "float_power(tf.constant(3.0),tf.constant(3.0))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0dhioJFVbyKJ"
      },
      "source": [
        "Remember that you can also pass tensors to `get_concrete_function`, in that case it returns the concrete function that would run for those inputs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "dzeSYHN1cGke"
      },
      "outputs": [],
      "source": [
        "row = tf.range(10)\n",
        "col = tf.constant([[1],[2],[3]])\n",
        "\n",
        "concrete_power = power.get_concrete_function(a = row, b = col)\n",
        "concrete_power(row, col).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JcsoDz8g9nJT"
      },
      "source": [
        "## Using a concrete function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "glxZkvRly1Am"
      },
      "source": [
        "A concrete function only accepts tensors as input:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1AIbwzSz0aCj"
      },
      "outputs": [],
      "source": [
        "float_power(tf.constant(2.0), tf.constant(3.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "7xq-82FokX_d"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  float_power(2.0,3.0)\n",
        "  assert False\n",
        "except ValueError:\n",
        "  traceback.print_exc(limit=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "V1eBDquxy5m5"
      },
      "source": [
        "It also only accepts inputs of the correct dtype:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "nSdGa_IJkayV"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  float_power(tf.constant(1),tf.constant(3))\n",
        "  assert False\n",
        "except tf.errors.InvalidArgumentError:\n",
        "  traceback.print_exc(limit=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kRc4YAxzsSdM"
      },
      "source": [
        "But it will try to execute even if the input tensors do not match the expected shape:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "MT-Z0cwrrrgZ"
      },
      "outputs": [],
      "source": [
        "float_power(tf.constant([1.,2.,3.,4.,5.]),tf.constant(3.)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "zWr690sEqwu0"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  float_power(tf.constant([1.,2.,3.]),tf.constant([4., 5.])).numpy()\n",
        "  assert False\n",
        "except tf.errors.InvalidArgumentError:  \n",
        "  traceback.print_exc(limit=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "PtHTfAtamxcu"
      },
      "source": [
        "By inspecting the concrete function you can see its inputs and outputs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "d-8tBcOHmw7x"
      },
      "outputs": [],
      "source": [
        "print(float_power.structured_input_signature)\n",
        "print(float_power.structured_outputs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "64z5eSmMzv8l"
      },
      "source": [
        "## Python Objects in signatures"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "1RTj0n2R0jGH"
      },
      "source": [
        "As you saw when tracing, each python object generates a new trace. Concrete functions represent a single `tf.Graph`, they don't do any retracing. When you call `get_concrete_function` with a python object as one of the arguments the object is **bound** to the function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "HxQtkTRYz-8u"
      },
      "outputs": [],
      "source": [
        "cube = power.get_concrete_function(\n",
        "    a = tf.TensorSpec([], dtype=tf.float32),\n",
        "    b = 3.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Iu8dRqCQ0SGW"
      },
      "source": [
        "This `cube` function no longer has a `b` argument:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "duRJIrDU0L-I"
      },
      "outputs": [],
      "source": [
        "print(cube.structured_input_signature)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "PnEQtxxA3OgX"
      },
      "outputs": [],
      "source": [
        "cube(tf.constant(10.0)).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hL76DCMVsdet"
      },
      "source": [
        "This is very similar to the way that standard python classes bind methods, and applies equally when you run `get_concrete_function` from a method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "dV2R0AOIsdJs"
      },
      "outputs": [],
      "source": [
        "class Greeter(object):\n",
        "  def __init__(self, greeting):\n",
        "    self.greeting = greeting\n",
        "\n",
        "  def greet(self, who):\n",
        "    return \" \".join([self.greeting, who])\n",
        "\n",
        "p = Greeter(\"Hello\")\n",
        "m = p.greet\n",
        "print(m)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "j9sEWo-wtP_K"
      },
      "outputs": [],
      "source": [
        "print(m(\"TensorFlow!\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "uv4X41_Htqof"
      },
      "source": [
        "When you have a `tf.function` decorating a method, similar rules apply:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "h131HNsj3nYm"
      },
      "outputs": [],
      "source": [
        "class MyModel(tf.Module):\n",
        "  def __init__(self, ins, outs):\n",
        "    initializer = tf.initializers.GlorotNormal()\n",
        "    self.W = tf.Variable(initializer([ins, outs]))\n",
        "    self.B = tf.Variable(tf.zeros([outs], dtype = tf.float32))\n",
        "\n",
        "  @tf.function\n",
        "  def run(self, x):\n",
        "    print('Tracing \"MyModule\":\\n    x={}\\n'.format(x))\n",
        "    return tf.matmul(x, self.W)+self.B"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "nGtPiUpZ3y0j"
      },
      "outputs": [],
      "source": [
        "mod = MyModel(ins=5, outs=3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3OVnEqJY6BHU"
      },
      "outputs": [],
      "source": [
        "mod.run([[1.0,1.0,1.0, 1.0, 1.0]]).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ca4tB9QguPBg"
      },
      "source": [
        "If you call the method's `.get_concrete_function`, the `self` is automatically bound as the first argument:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "0buNZlUw6XNN"
      },
      "outputs": [],
      "source": [
        "concrete_run = mod.run.get_concrete_function(x = tf.TensorSpec([None, None]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "D-v1YUN562wH"
      },
      "outputs": [],
      "source": [
        "concrete_run(tf.constant([[1.0,1.0,1.0, 1.0, 1.0],\n",
        "                          [2.0,2.0,2.0, 2.0, 2.0]])).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "voRoyyNBuq5-"
      },
      "source": [
        "See how `self` is no longer part of the input signature:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "t4V1ChR97uNS"
      },
      "outputs": [],
      "source": [
        "print(concrete_run.structured_input_signature)\n",
        "print(concrete_run.structured_outputs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "WlHhcO30mZYN"
      },
      "source": [
        "## Accessing concrete functions from a SavedModel\n",
        "\n",
        "When you save a [SavedModel](saved_model.ipynb) you're really saving the `tf.function's` cache of concrete functions.\n",
        "\n",
        "Because concrete functions are generated by tracing the input you need to execute at least one trace to save a SavedModel.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "L5iunDSWxS3b"
      },
      "outputs": [],
      "source": [
        "dropout = Dropout(0.5)\n",
        "\n",
        "_ = dropout(tf.range(10, dtype=tf.float32), tf.constant(True))\n",
        "_ = dropout(tf.random.normal([2, 3]), tf.constant(True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "QMr2Bz0FQeLN"
      },
      "source": [
        "Note: `tf.saved_model` retraces all `concrete_functions` when saving them. This is to ensure that the exported concrete functions capture changes in the environment on export (e.g. distribution strategy scope)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "OtoZNlykta1L"
      },
      "outputs": [],
      "source": [
        "export_dir = 'dropout'\n",
        "tf.saved_model.save(dropout, export_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "83oSBW8jxau1"
      },
      "source": [
        "### Direct access\n",
        "\n",
        "When you load a `tf.saved_model` your methods are restored as polymorphic functions:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "o8fVspUQtgZs"
      },
      "outputs": [],
      "source": [
        "reloaded_dropout = tf.saved_model.load(export_dir)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "4IB6iNTUxLpk"
      },
      "outputs": [],
      "source": [
        "print(reloaded_dropout(tf.range(10, dtype=tf.float32), tf.constant(False)).numpy())\n",
        "print(reloaded_dropout(tf.random.normal([2,3]), tf.constant(True)).numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "237FYOTTxKeI"
      },
      "source": [
        "But since the `saved_model` only contains the cache of concrete functions (an d not the python source and data), it cannot handle signatures that don't match: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "NKIE3TxLwzj8"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  reloaded_dropout(tf.range(12, dtype=tf.float32), tf.constant(True))\n",
        "  assert False\n",
        "except ValueError:\n",
        "  traceback.print_exc(limit=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Fy8Bryv5yHFs"
      },
      "source": [
        "From the reloaded module you can select a specific concrete function instead of relying on the dispatch by, again, using the `get_concrete_function` method: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mxr0uk9YyGQu"
      },
      "outputs": [],
      "source": [
        "cf = reloaded_dropout.__call__.get_concrete_function(\n",
        "    x = tf.TensorSpec([10]), \n",
        "    training = tf.TensorSpec([], tf.bool))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1jRI6tByyX_8"
      },
      "outputs": [],
      "source": [
        "result = cf(tf.range(10, dtype=tf.float32), tf.constant(True)).numpy()\n",
        "print(result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "BCM5i0pZ2Vku"
      },
      "source": [
        "### Named signatures: Exporting for C++"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-akdpHabyEIy"
      },
      "source": [
        "C++ consumers of SavedModels do not use the above \"Direct Access\" method, or it's dynamic dispatch, to get and run concrete functions from the SavedModel. \n",
        "\n",
        "They use a more explicit interface called \"exported signatures\", where you specify exactly which concrete functions to export. \n",
        "\n",
        "You specify the concrete functions to export by passing a `signatures` argument to `tf.saved_model.save`.\n",
        "\n",
        "It takes either:\n",
        "\n",
        "* A dictionary of functions. This allows you to name each function.\n",
        "* A single function. When a single function is exported, it will be named \"serving_default\", using `saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY`.\n",
        "\n",
        "These signatures are required when using TensorFlow Serving.\n",
        "\n",
        "Note: These \"exported signatures\" are wraped to always return a dictionary of results."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "LeBeEh8t3j9Z"
      },
      "source": [
        "#### Simple example"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "8ji4CDMVzTQn"
      },
      "outputs": [],
      "source": [
        "dropout = Dropout(0.5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Sfge8hwHyJOl"
      },
      "outputs": [],
      "source": [
        "cf = dropout.__call__.get_concrete_function(tf.zeros((2,3), dtype=tf.float32), tf.constant(False))\n",
        "\n",
        "import time\n",
        "export_dir = \"./saved/\"+str(time.time())\n",
        "\n",
        "tf.saved_model.save(dropout, export_dir, signatures = cf)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "6pP_EgrY2Ces"
      },
      "source": [
        "This `saved_model` only contains the one signature, and it can be recovered by name, from the `signatures` dictionary:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Pt6Z8ZUH2Dc8"
      },
      "outputs": [],
      "source": [
        "reloaded = tf.saved_model.load(export_dir)\n",
        "\n",
        "print(reloaded.signatures)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ODMPWghD4IRC"
      },
      "source": [
        "When using a \"exported signatures\" these concrete functions always return a dictionary of outputs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "USuqYJY84F4F"
      },
      "outputs": [],
      "source": [
        "cf = reloaded.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]\n",
        "result = cf(x=tf.random.normal([2,3]), training=tf.constant(True))\n",
        "\n",
        "print(result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "tlnJz2AEyyHV"
      },
      "source": [
        "In the example above, the output names auto-generated by the signature is fairly generic. You can check the output names using the `structured_outputs` method:\n",
        "\n",
        "You can check the expected output-tensor names using the `.structured_outputs` method: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "bJKjvia-xS-_"
      },
      "outputs": [],
      "source": [
        "cf.structured_outputs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "EAKTLuvUrg21"
      },
      "source": [
        "Typically you wannt to set the output names yourself."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "g7y03GXg3pii"
      },
      "source": [
        "#### Example: Setting the output names"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "LEkFHgo81Iex"
      },
      "source": [
        "To control the names of the outputs, modify your `tf.function` to return a dictionary that maps names to output tensors.:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "kz8QOgr21VZR"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def named_result(x, training=True):\n",
        "  return {'dropout': dropout(x, training)}\n",
        "\n",
        "dropout.named_result = named_result\n",
        "\n",
        "cf = dropout.named_result.get_concrete_function(tf.zeros((2,3), dtype=tf.float32),\n",
        "                                                tf.constant(False))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Nhg_2t-Jr-Pn"
      },
      "source": [
        "#### Example: Setting the signature names"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "szEQyfU83yQ0"
      },
      "source": [
        "To set the name of the signature pass a dictionary of concrete functions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "fUL6ntYu31TL"
      },
      "outputs": [],
      "source": [
        "export_dir = \"./saved/\"+str(time.time())\n",
        "tf.saved_model.save(dropout, export_dir, signatures = {'simple':cf})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aFw4kDvU195u"
      },
      "outputs": [],
      "source": [
        "reloaded = tf.saved_model.load(export_dir)\n",
        "cf = reloaded.signatures['simple']\n",
        "result = cf(x=tf.random.normal([2,3]), training=tf.constant(True))\n",
        "\n",
        "print({key:value.numpy() for key,value in result.items()})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Ho7_3_G13jVF"
      },
      "source": [
        "To specify multiple signatures pass a dictionary of `(name, concrete_function)` pairs to `saved_model.save`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Z0vgRTcE2pvi"
      },
      "outputs": [],
      "source": [
        "vector = dropout.__call__.get_concrete_function(tf.TensorSpec((2,3), dtype=tf.float32), tf.constant(False))\n",
        "matrix = dropout.__call__.get_concrete_function(tf.TensorSpec((2,3), dtype=tf.float32), tf.constant(False))\n",
        "cube = dropout.__call__.get_concrete_function(tf.TensorSpec((2,3), dtype=tf.float32), tf.constant(False))\n",
        "\n",
        "export_dir = \"./saved/\"+str(time.time())\n",
        "\n",
        "tf.saved_model.save(dropout, export_dir, \n",
        "                    signatures = {\n",
        "                        \"vector\": vector,\n",
        "                        \"matrix\": matrix,\n",
        "                        \"cube\": cube\n",
        "                    })"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "L7jpWSND4_VY"
      },
      "source": [
        "Now reload that model and inspect the signature listing:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "YTVaGx4W42EO"
      },
      "outputs": [],
      "source": [
        "reloaded = tf.saved_model.load(export_dir)\n",
        "print('{}'.format(reloaded.signatures).replace(\"{\",\"{\\n    \").replace(\"\u003e, \", \"\u003e,\\n    \"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "cbVqU7xOsLR3"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "concrete_function.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
